/*
 *         (C) COPYRIGHT NBI Limited
 *              ALL RIGHT RESERVED
 *
 * File        : fastflow/graph/node.h
 * Authors     : dzhang
 * Create Time : 2021/09/04 18:10:23
 * Description :
 *
 */

#ifndef __FASTFLOW_GRAPH_NODE_H___
#define __FASTFLOW_GRAPH_NODE_H___

#include <string>

#include "fastflow/core/common.h"
#include "fastflow/graph/node_arg.h"
#include "fastflow/graph/const_pointer_container.h"
#include "fastflow/graph/function.h"

namespace fastflow {

class Graph;
class Node;

/**
@class Node
Class representing a node in the graph.
*/
class Node {
public:
    /** Node types */
    enum class Type {
        Primitive = 0,  ///< The node refers to a primitive operator.
        Fused = 1,      ///< The node refers to a function.
    };

    ~Node() = default;

    /**
    @class EdgeEnd
    Class representing the end of an edge. It could be an input or output edge end of a node.
    For the node's input edge end, it's the source end, as the destination end is the node itself.
    For the node's output edge end, it's the destination end, as the source end is the node itself.
    */
    class EdgeEnd {
    public:
        /**
        Construct an EdgeEnd
        @param node The source node if this is an input edge to the current node,
        or the destination node if this is an output edge from the current node.
        @param node_arg The NodeArg to use for the edge.
        */
        EdgeEnd(const Node& node, const NodeArg& node_arg) noexcept;

        /** Construct a control edge.
        @param node The node the edge joins to the current node.
        */
        EdgeEnd(const Node& node) noexcept;

        /** Gets the Node that this EdgeEnd refers to. */
        const Node& GetNode() const noexcept;

        /** Gets the NodeArg that this edge end refers to.
        @returns NodeArg pointer or nullptr if this is a control edge. */
        const NodeArg* GetNodeArg() const noexcept;

    private:
        const Node* node_;
        const NodeArg* node_arg_;
    };

    /** Gets the Node's NodeIndex. */
    NodeIndex Index() const noexcept;

    /** Gets the Node's name. */
    const std::string& Name() const noexcept;

    /** Gets the Node's operator type. */
    const std::string& OpType() const noexcept;

    /** Gets the domain of the OperatorSet that specifies the operator returned by #OpType. */
    const std::string& Domain() const noexcept;

    /** Gets the opset version that the Node's operator was first defined in.
    @returns Opset version. If -1 the Node's operator has not been set.
    @remarks Prefer over Op()->SinceVersion() as Op() is disabled in a minimal build
    */
    int SinceVersion() const noexcept { return since_version_; }


    /** Gets the Node's OpSchema.
    @remarks The graph containing this node must be resolved, otherwise nullptr will be returned. */
    const ONNX_NAMESPACE::OpSchema* Op() const noexcept;

    /** Gets the Node's Node::Type. */
    Node::Type NodeType() const noexcept;

    /** Gets the function body if the #NodeType is fused, or nullptr if not. */
    const Function* GetFunctionBody() const noexcept;

    /** Gets the node description. */
    const std::string& Description() const noexcept;

    /**
    Helper to iterate through the container returned by #InputDefs() or #OutputDefs() and call the provided function.
    @param node_args Collection of NodeArgs returned by #InputDefs() or #OutputDefs()
    @param func Function to call for each valid NodeArg in the node_args. The function is called with the NodeArg
                and the index number in the container.
    @returns Status with success or error information.
    @remarks Returns immediately on error.
    */
    static Status ForEachWithIndex(const ConstPointerContainer<std::vector<NodeArg*>>& node_args,
                                           std::function<Status(const NodeArg& arg, size_t index)> func) {
        for (size_t index = 0; index < node_args.size(); ++index) {
            auto arg = node_args[index];
            if (!arg->Exists())
                continue;
            FASTFLOW_RETURN_IF_ERROR(func(*arg, index));
        }
        return Status::OK();
    }

    /** Gets the Node's input definitions.
    @remarks requires ConstPointerContainer wrapper to apply const to the NodeArg pointers so access is read-only. */
    const ConstPointerContainer<std::vector<NodeArg*>> InputDefs() const noexcept {
        return ConstPointerContainer<std::vector<NodeArg*>>(definitions_.input_defs);
    }

    /** Gets a modifiable collection of the Node's input definitions. */
    std::vector<NodeArg*>& MutableInputDefs() noexcept {
        return definitions_.input_defs;
    }

    /** Gets the count of arguments for each of the Node's explicit inputs. */
    const std::vector<int>& InputArgCount() const noexcept {
        return definitions_.input_arg_count;
    }

    /** Gets a modifiable count of arguments for each of the Node's explicit inputs.
    @todo This should be removed in favor of a method that updates the input args and the count.
          Currently these operations are separate which is not a good setup. */
    std::vector<int>& MutableInputArgsCount() {
        return definitions_.input_arg_count;
    }

    /** Gets the implicit inputs to this Node.
    If this Node contains a subgraph, these are the NodeArg's that are implicitly consumed by Nodes within that
    subgraph. e.g. If and Loop operators.*/
    const std::vector<const NodeArg*>& ImplicitInputDefs() const noexcept {
        return definitions_.implicit_input_defs;
    }

    /** Gets the Node's output definitions.
    @remarks requires ConstPointerContainer wrapper to apply const to the NodeArg pointers so access is read-only. */
    const ConstPointerContainer<std::vector<NodeArg*>> OutputDefs() const noexcept {
        return ConstPointerContainer<std::vector<NodeArg*>>(definitions_.output_defs);
    }

    /** Struct to provide sorting between EdgeEnd instances based on NodeIndex first, and NodeArg::Name second. */
    struct EdgeEndCompare {
        bool operator()(const EdgeEnd& lhs, const EdgeEnd& rhs) const {
            if (lhs.GetNode().Index() == rhs.GetNode().Index()) {
                auto lhs_arg = lhs.GetNodeArg();
                auto rhs_arg = rhs.GetNodeArg();
                std::string lhs_arg_name = lhs_arg == nullptr ? "" : lhs_arg->Name();
                std::string rhs_arg_name = rhs_arg == nullptr ? "" : rhs_arg->Name();
                return lhs_arg_name.compare(rhs_arg_name) < 0;
            }
            return lhs.GetNode().Index() < rhs.GetNode().Index();
        }
    };

    using EdgeSet = std::set<EdgeEnd, EdgeEndCompare>;
    using EdgeConstIterator = EdgeSet::const_iterator;

    /**
    @class NodeConstIterator
    Class to provide const access to Node instances iterated via an EdgeConstIterator. */
    class NodeConstIterator {
    public:
        NodeConstIterator(EdgeConstIterator p_iter);

        bool operator==(const NodeConstIterator& p_other) const;

        bool operator!=(const NodeConstIterator& p_other) const;

        void operator++();
        void operator--();

        const Node& operator*();

    private:
        EdgeConstIterator m_iter;
    };

    // Functions defined to traverse a Graph as below.

    /** Gets an iterator to the beginning of the input nodes to this Node. */
    NodeConstIterator InputNodesBegin() const noexcept {
        return NodeConstIterator(relationships_.input_edges.cbegin());
    };
    /** Gets an iterator to the end of the input nodes to this Node. */
    NodeConstIterator InputNodesEnd() const noexcept {
        return NodeConstIterator(relationships_.input_edges.cend());
    }

    /** Gets an iterator to the beginning of the output nodes from this Node. */
    NodeConstIterator OutputNodesBegin() const noexcept {
        return NodeConstIterator(relationships_.output_edges.cbegin());
    }
    /** Gets an iterator to the end of the output nodes from this Node. */
    NodeConstIterator OutputNodesEnd() const noexcept {
        return NodeConstIterator(relationships_.output_edges.cend());
    }

    /** Gets an iterator to the beginning of the input edges to this Node.
    @remarks There are no nullptr entries in this collection. */
    EdgeConstIterator InputEdgesBegin() const noexcept {
        return relationships_.input_edges.cbegin();
    }

    /** Gets an iterator to the end of the input edges to this Node. */
    EdgeConstIterator InputEdgesEnd() const noexcept {
        return relationships_.input_edges.cend();
    }

    /** Gets an iterator to the beginning of the output edges from this Node.
    @remarks There are no nullptr entries in this collection. */
    EdgeConstIterator OutputEdgesBegin() const noexcept {
        return relationships_.output_edges.cbegin();
    }

    /** Gets an iterator to the end of the output edges from this Node. */
    EdgeConstIterator OutputEdgesEnd() const noexcept {
        return relationships_.output_edges.cend();
    }

    /** Gets the Node's control inputs. */
    const std::set<std::string>& ControlInputs() const noexcept {
        return relationships_.control_inputs;
    }

    /** Gets the number of input edges to this Node */
    size_t GetInputEdgesCount() const noexcept {
        return relationships_.input_edges.size();
    }

    /** Gets the number of output edges from this Node */
    size_t GetOutputEdgesCount() const noexcept {
        return relationships_.output_edges.size();
    }

    /** Add an attribute to this Node with specified attribute name and value. */
    void AddAttribute(const std::string& attr_name, const ONNX_NAMESPACE::AttributeProto& value);

#define ADD_ATTR_INTERFACES(TypeName)                                                                                  \
    void AddAttribute(const std::string& attr_name, const TypeName& value);                                            \
    void AddAttribute(const std::string& attr_name, const std::vector<TypeName>& values);

    ADD_ATTR_INTERFACES(int64_t)
    ADD_ATTR_INTERFACES(float)
    ADD_ATTR_INTERFACES(std::string)
    ADD_ATTR_INTERFACES(ONNX_NAMESPACE::TensorProto)
    ADD_ATTR_INTERFACES(ONNX_NAMESPACE::GraphProto)

    /** Remove the specified attribute from this Node */
    bool ClearAttribute(const std::string& attr_name);

    /** Gets the Node's attributes. */
    const NodeAttributes& GetAttributes() const noexcept;

    /** Gets the Graph instance that is instantiated from a GraphProto attribute during Graph::Resolve.
    @param attr_name Attribute name for the GraphProto attribute.
    @returns nullptr if the Graph instance has not been instantiated or attribute does not contain a GraphProto.
    */
    const Graph* GetGraphAttribute(const std::string& attr_name) const;

    /** Gets the mutable Graph instance that is instantiated from a GraphProto attribute during Graph::Resolve.
    @param attr_name Attribute name for the GraphProto attribute.
    @returns nullptr if the Graph instance has not been instantiated or attribute does not contain a GraphProto.
    */
    Graph* GetMutableGraphAttribute(const std::string& attr_name);

    /** Gets the execution ProviderType that this node will be executed by. */
    ProviderType GetExecutionProviderType() const noexcept;

    /** Sets the execution ProviderType that this Node will be executed by. */
    void SetExecutionProviderType(ProviderType execution_provider_type);

    /** Gets the NodeProto representation of this Node. */
    void ToProto(ONNX_NAMESPACE::NodeProto& proto) const;

    /** Call the provided function for all explicit inputs, implicit inputs, and outputs of this Node.
        If the NodeArg is an explicit or implicit input, is_input will be true when func is called. */
    void ForEachDef(std::function<void(const fastflow::NodeArg&, bool is_input)> func) const;

    /** Replaces any matching definitions in the Node's explicit inputs or explicit outputs.
    @param replacements Map of current NodeArg to replacement NodeArg.
    */
    void ReplaceDefs(const std::map<const fastflow::NodeArg*, fastflow::NodeArg*>& replacements);

    /**
    @class Definitions
    The input and output definitions for this Node.
    */
    class Definitions {
    public:
        Definitions() noexcept = default;

        /** The Node's explicit input definitions. */
        std::vector<NodeArg*> input_defs;

        /**
        The number of inputs for each argument of the operator or function which this node refers.
        @remarks For example, #input_defs has 10 elements (inputs), and #input_arg_count is {4, 6}.
        This means that 4 elements (inputs) of input_defs map to the first argument of the operator or function, and
        the other 6 map to the second argument.
        */
        std::vector<int> input_arg_count;

        /** The Node's output definitions. */
        std::vector<NodeArg*> output_defs;

        /** The Node's implicit input definitions if the Node contains one or more subgraphs
        (i.e. GraphProto attributes) and the subgraph/s implicitly consume these values.
        @remarks For example, a subgraph in an 'If' node gets all its input values via this mechanism rather than
        there being explicit inputs to the 'If' node that are passed to the subgraph.
        They are pseudo-inputs to this Node as it has an implicit dependency on them. */
        std::vector<const NodeArg*> implicit_input_defs;

    private:
        FASTFLOW_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(Definitions);
    };

    /**
    @class Relationships
    Defines the relationships between this Node and other Nodes in the Graph.
    */
    class Relationships {
    public:
        Relationships() = default;

        void Clear() noexcept {
            input_edges.clear();
            output_edges.clear();
            control_inputs.clear();
        }

        /** The edges for Nodes that provide inputs to this Node. */
        EdgeSet input_edges;

        /** The edges for Nodes that receive outputs from this Node. */
        EdgeSet output_edges;

        /** The Node names of the control inputs to this Node. */
        std::set<std::string> control_inputs;

    private:
        FASTFLOW_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(Relationships);
    };

private:
    FASTFLOW_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(Node);

    // NOTE: This friendship relationship should ONLY be used for calling methods of the Node class and not accessing
    // the data members directly, so that the Node can maintain its internal invariants.
    friend class Graph;

    Node(NodeIndex index, Graph& graph) : index_(index), graph_(&graph) {}

    void Init(const std::string& name, const std::string& op_type, const std::string& description,
              const std::vector<NodeArg*>& input_args, const std::vector<NodeArg*>& output_args,
              const NodeAttributes* attributes, const std::string& domain);

    // create a Graph instance for an attribute that contains a GraphProto
    void CreateSubgraph(const std::string& attr_name);

    // internal only method to allow selected classes to directly alter the input/output definitions and arg counts
    Definitions& MutableDefinitions() noexcept;

    // internal only method to allow selected classes to directly alter the links between nodes.
    Relationships& MutableRelationships() noexcept;

    const std::vector<std::unique_ptr<Graph>>& MutableSubgraphs() noexcept {
        return subgraphs_;
    }

    const Definitions& GetDefinitions() const noexcept {
        return definitions_;
    }
    const Relationships& GetRelationships() const noexcept {
        return relationships_;
    }

    void SetNodeType(Node::Type node_type) noexcept;

    void SetFunctionBody(const Function& func);

    // validate and update the input arg count
    Status UpdateInputArgCount();

    // Node index. Default to impossible value rather than 0.
    NodeIndex index_ = std::numeric_limits<NodeIndex>::max();

    // Node name.
    std::string name_;

    // Node operator type.
    std::string op_type_;

    // OperatorSet domain of op_type_.
    std::string domain_;

    // OperatorSchema that <*this> node refers to.
    const ONNX_NAMESPACE::OpSchema* op_ = nullptr;
    Node::Type node_type_ = Node::Type::Primitive;

    // The function body is owned by graph_
    const Function* func_body_ = nullptr;

    // Node doc string.
    std::string description_;

    // input/output defs and arg count
    Definitions definitions_;

    // set from op_->SinceVersion() or via deserialization when OpSchema is not available
    int since_version_ = -1;

    // Relationships between this node and others in the graph
    Relationships relationships_;

    // Device.
    std::string execution_provider_type_;

    // Map from attribute name to attribute.
    // This allows attribute adding and removing.
    NodeAttributes attributes_;

    // Graph that contains this Node
    Graph* graph_;

    // Map of attribute name to the Graph instance created from the GraphProto attribute
    std::unordered_map<std::string, Graph*> attr_to_subgraph_map_;

    // Graph instances for subgraphs that are owned by this Node
    std::vector<std::unique_ptr<Graph>> subgraphs_;
};

}  // namespace fastflow

#endif  // __FASTFLOW_GRAPH_NODE_H___
